import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation

# 定义网格的大小
grid_size = 100

# 定义起点和目标点的位置
start = (0, 0)
goal = [99, 99]

# 定义运动方向
actions = ['up', 'down', 'left', 'right', 'upleft', 'upright', 'downleft', 'downright']
num_actions = len(actions)

# 定义Q-table，并初始化为0
Q = np.zeros((grid_size, grid_size, num_actions))

# 定义参数
alpha = 0.1  # 学习率
gamma = 0.9  # 折扣因子
epsilon = 0.1  # 探索率
num_episodes = 5000
path = []
steps = 0


# 定义辅助函数，将位置转换为状态
def state_from_position(position):
    return position[0], position[1]


# 定义辅助函数，将状态转换为位置
def position_from_state(state):
    return state[0], state[1]


# 定义辅助函数，选择动作
def choose_action(state):
    if np.random.uniform() < epsilon or np.sum(Q[state]) == 0:
        return np.random.randint(num_actions)
    else:
        return np.argmax(Q[state])


# 执行动作
def executeAction(action, position):
    next_position = list(position)

    if actions[action] == 'up':
        next_position[0] = max(position[0] - 1, 0)
    elif actions[action] == 'down':
        next_position[0] = min(position[0] + 1, grid_size - 1)
    elif actions[action] == 'left':
        next_position[1] = max(position[1] - 1, 0)
    elif actions[action] == 'right':
        next_position[1] = min(position[1] + 1, grid_size - 1)
    elif actions[action] == 'upleft':
        next_position[0] = max(position[0] - 1, 0)
        next_position[1] = max(position[1] - 1, 0)
    elif actions[action] == 'upright':
        next_position[0] = max(position[0] - 1, 0)
        next_position[1] = min(position[1] + 1, grid_size - 1)
    elif actions[action] == 'downleft':
        next_position[0] = min(position[0] + 1, grid_size - 1)
        next_position[1] = max(position[1] - 1, 0)
    elif actions[action] == 'downright':
        next_position[0] = min(position[0] + 1, grid_size - 1)
        next_position[1] = min(position[1] + 1, grid_size - 1)

    return next_position


minLen = 1000000000
tempPath = []
# 开始训练
print('Training...')
for episode in range(num_episodes):
    tempPath = []
    if episode % 1000 == 0:
        print('Episode:', episode + 1, '/', num_episodes, ";  本次训练步数：", steps)
    position = start
    state = state_from_position(position)
    steps = 0
    if episode == num_episodes - 1:
        gamma = 0.95
    while position != goal:
        steps += 1
        tempPath.append(position)
        action = choose_action(state)
        # next_position = list(position)

        next_position = executeAction(action, position)

        next_state = state_from_position(next_position)
        # 计算欧式距离
        reward = -1
        # if (abs(position[0] - 99) ** 2 + abs(position[1] - 99) ** 2) != 0:
        #     reward = 500 / (abs(position[0] - 99) ** 2 + abs(position[1] - 99) ** 2)

        Q[state][action] += alpha * (reward + gamma * np.max(Q[next_state]) - Q[state][action])

        position = next_position
        state = next_state

    if steps < minLen:
        minLen = steps
        path = tempPath.copy()

# 从起点开始，根据Q-table选择最优路径
# position = start
# state = state_from_position(position)
# path = [position]
#
# priority = 1
# while position != (99, 99):
#     if priority > len(actions):
#         print('No path found!')
#         break
#     action = np.argmax(Q[state]) #argmax()返回的是数组值从小到大的索引值
#     # 选择Q中第priority大的值 argsort()函数返回的是数组值从小到大的索引值
#     next_position = executeAction(action, position)
#
#     # 如果没有更新位置，则选择下一个动作
#     if next_position == position:
#         priority += 1
#     else:
#         priority = 1
#
#     position = next_position
#     state = state_from_position(position)
#     path.append(position)

# 可视化最短路径
path = np.array(path)
# plt.plot(path[:, 1], path[:, 0], linewidth=2)
line = plt.plot(path[:, 1], path[:, 0], linewidth=2)


def init():
    line[0].set_data([], [])
    return line


def animate(i):
    line[0].set_data(path[:i, 1], path[:i, 0])
    return line


# plt.plot(path[:, 1], path[:, 0], linewidth=2)
ani = animation.FuncAnimation(plt.figure(), animate, init_func=init, frames=len(path), interval=100)

plt.xlim(0, grid_size - 1)
plt.ylim(0, grid_size - 1)
plt.xticks(np.arange(0, grid_size, 10))
plt.yticks(np.arange(0, grid_size, 10))
plt.grid(True)
plt.xlabel('Column')
plt.ylabel('Row')
plt.title('Shortest Path')
plt.show()
